#include <iostream>
#include <fstream>
#include "FeatureVectors.h"
#include "SVMClassificationModule.h"
#include "svm.h"
using namespace std;
/*******************************************************************/
CSVMClassificationModule::CSVMClassificationModule()
{
	m_classNameEXT = "txt";
	m_svmNameEXT = "svm";
	m_extractor = NULL;
	m_svmModel = NULL; 
}
/*******************************************************************/
CSVMClassificationModule::~CSVMClassificationModule()
{
	if( NULL != m_svmModel )
	{
		svm_free_and_destroy_model( &m_svmModel ); 
	}
}
/*******************************************************************/
bool CSVMClassificationModule::Load( string fname, string path )
{
	bool retVal = true; 
	retVal &= LoadSVMFile( fname, path ); 
	retVal &= LoadClassNameFile( fname, path );
	retVal &= m_normalizer.Load( fname, path ); 
	return retVal; 
}
/*******************************************************************/
bool CSVMClassificationModule::Save( string fname, string path )
{
	bool retVal = false; 
	if( m_svmModel != NULL )
	{
		retVal = true; 
		retVal &= SaveSVMFile( fname, path ); 
		retVal &= SaveClassNameFile( fname, path );
		retVal &= m_normalizer.Save( fname, path ); 
	}
	return retVal; 
}
/*******************************************************************/
void CSVMClassificationModule::SetClassLabels( TClassToNameMap& classLabels )
{
	m_ClassNameMap = classLabels;
}
/*******************************************************************/
TClassToNameMap CSVMClassificationModule::GetClassLabels( void )
{
	return m_ClassNameMap;
}
/*******************************************************************/
string CSVMClassificationModule::GetClassName( double value ) 
{
	return m_ClassNameMap[value];
}
/*******************************************************************/
void CSVMClassificationModule::SetFeatureExtractor( CFeatureExtractorBase * extractor ) 
{
	m_extractor = extractor;
}
/*******************************************************************/
CFeatureExtractorBase* CSVMClassificationModule::GetFeatureExtractor( void )
{
	return m_extractor; 
}
/*******************************************************************/  
void CSVMClassificationModule::SetParams( svm_parameter params ) 
{
	m_parameters = params; 
}
/*******************************************************************/
svm_parameter CSVMClassificationModule::GetParams( void ) 
{
	return m_parameters; 
}
/*******************************************************************/
bool CSVMClassificationModule::BuildSVM( CDataPointManager * data, bool manageImgs  )
{
	bool retVal = false;
	if( NULL != m_svmModel )
	{
		svm_free_and_destroy_model( &m_svmModel ); 
	}
	TDataSet feat_data = m_extractor->ExtractFeaturesFromData(data,manageImgs);
	TDataSet ndata = m_normalizer.CalculateNormalizationValues(feat_data);
	int samples = static_cast<int>(ndata.size());
	int realFeats = static_cast<int>(ndata.begin()->size());
	int nfeats = realFeats+1; // SVM's feature space appears to be plus one for some sort of terminator
	m_problem.l = samples; // size??
	m_problem.y = new double[samples]; // thinking this is label; 
	m_problem.x = new svm_node *[samples];
	svm_node *feature_space = new svm_node[nfeats * samples];
	TDataSetIter iter;
	double label=0.00;
	int i=0;
	for(  iter = ndata.begin(); iter != ndata.end(); ++iter, ++i)
	{ 
		for( int j=0;j<realFeats;++j)
		{
			feature_space[nfeats*i+j].value = iter->at(j);
			feature_space[nfeats*i+j].index = j+1;
		}
		feature_space[nfeats*i+realFeats].index = -1;
		m_problem.x[i]=&feature_space[nfeats*i];
		label = data->GetDataPoint(i)->GetClassification().GetTruth();
		m_problem.y[i] = static_cast<int>(label);
	}

	m_svmModel = svm_train(&m_problem, &m_parameters);
	if( NULL != m_svmModel )
	{
		retVal = true; 
	}
	return retVal; 
}
/*******************************************************************/ 
bool CSVMClassificationModule::BuildSVM( TDataMgrVector& data, TTruthLabels classes, bool manageImgs )
{
	bool retVal = false;
	if( NULL != m_svmModel )
	{
		svm_free_and_destroy_model( &m_svmModel ); 
	}
	int dmSz = static_cast<int>(data.size());
	TTrainingDataSet dataSet;
	TDataSet bigSet; 
	for( int i=0; i < dmSz; i++ )
	{
		TDataSet temp = m_extractor->ExtractFeaturesFromData(data[i],manageImgs);
		dataSet[classes[i]]= temp; 
		bigSet.insert(bigSet.begin(), temp.begin(), temp.end() ); 
	}
	m_normalizer.CalculateNormalizationValues(bigSet); 
	bigSet.clear(); 
	TTrainingDataSetIter tdsIter; 
	// do the normalization
	for( tdsIter = dataSet.begin(); tdsIter!= dataSet.end(); tdsIter++ )
	{
		m_normalizer.Normalize(tdsIter->second);
	}

	// get the total samples so we can allocate
	TTrainingDataSetIter classIter;
	int samples = 0; 
	for( classIter = dataSet.begin(); classIter != dataSet.end(); ++classIter)
	{
		samples += static_cast<int>(((TDataSet)classIter->second).size());
	}

	TDataSet tempDataSet = dataSet.begin()->second;
	int realFeats = static_cast<int>(tempDataSet.begin()->size());
	int nfeats = realFeats+1; // SVM's feature space appears to be plus one for some sort of terminator
	m_problem.l = samples; // size??
	m_problem.y = new double[samples]; // thinking this is label; 
	m_problem.x = new svm_node *[samples];
	svm_node *feature_space = new svm_node[nfeats * samples];
	TDataSetIter iter;
	int i=0;
	for(  classIter = dataSet.begin(); classIter != dataSet.end(); ++classIter )
	{
		tempDataSet = classIter->second;
		double label = classIter->first; 
		TDataSetIter iter; 
		for( iter = tempDataSet.begin(); iter != tempDataSet.end(); ++iter, ++i)
		{
			for( int j=0;j<realFeats;++j)
			{
				feature_space[nfeats*i+j].value = iter->at(j);
				feature_space[nfeats*i+j].index = j+1;
			}
			feature_space[nfeats*i+realFeats].index = -1;
			m_problem.x[i]=&feature_space[nfeats*i];
			m_problem.y[i] = static_cast<int>(label);
		}
	}

	m_svmModel = svm_train(&m_problem, &m_parameters);
	if( NULL != m_svmModel )
	{
		retVal = true; 
	}
	return retVal;  
}
/*******************************************************************/
double CSVMClassificationModule::Classify( IplImage* image )
{
	double retVal = 0.00; 
	TFeatureVector data = m_extractor->ExtractFeaturesFromImage(image);
	m_normalizer.Normalize(data);
	int realFeats = static_cast<int>(data.size()); 
	int nfeats = realFeats+1; 
	svm_node* test_value = new svm_node[nfeats];
	for( int j=0;j<realFeats;++j)
	{
		test_value[j].value = data[j];
		test_value[j].index = j+1;
	}
	test_value[nfeats].index = -1;

	retVal = svm_predict(m_svmModel, test_value);
	//cout << "Classification: " << retVal << endl; 
	/*if( NULL != test_value )
	{
		delete  test_value; 
	}*/
	return retVal;
}
/*******************************************************************/
double CSVMClassificationModule::Classify( CTruthDataPoint* datapoint, bool manageImgs )
{
	double retVal = 0.00;
	if( manageImgs )
	{
		datapoint->LoadSourceImage();	
	}
	if( datapoint->HasImage() )
	{
		retVal = Classify( datapoint->GetImage() ); 
		datapoint->GetClassification().SetExperiment(retVal);
		if( manageImgs )
		{
			datapoint->FreeImage();
		}
	}
	return retVal;
}
/*******************************************************************/
TTruthLabels CSVMClassificationModule::Classify( CDataPointManager * data, bool manageImgs ) 
{
	TTruthLabels retVal;
	double classification = 0.00;
	for( int i =0; i < data->GetSize(); i++)
	{
		classification = Classify( data->GetDataPoint(i),manageImgs);
		retVal.push_back(classification); 
	}
	return retVal; 
}
/*******************************************************************/
bool CSVMClassificationModule::LoadSVMFile(string fname, string path)
{
	bool retVal = false;
	if( NULL != m_svmModel )
	{
		svm_free_and_destroy_model( &m_svmModel ); 
	}
	string fullName = path + fname +"." + m_svmNameEXT;
	m_svmModel = svm_load_model(fullName.c_str());
	if( m_svmModel != NULL )
	{
		retVal = true; 
	}
	return retVal; 
}
/*******************************************************************/
bool CSVMClassificationModule::LoadClassNameFile( string fname, string path )
{
	bool retVal = false; 
	string fullName = path + fname + "." + m_classNameEXT;
	ifstream ifs;
	ifs.open(fullName.c_str() );
	if( ifs.is_open() )
	{
		int sz=0;
		m_ClassNameMap.clear(); 
		string name = ""; 
		double val = 0.00;
		ifs >> sz;
		for( int i=0; i < sz; i++ )
		{
			ifs >> val >> name; 
			m_ClassNameMap[val] = name;
		}
		if( ifs.good() )
		{
			retVal = true;
		}
		ifs.close();
	}
	return retVal; 
}
/*******************************************************************/
bool CSVMClassificationModule::SaveSVMFile(string fname, string path)
{
	bool retVal = false; 
	if( NULL != m_svmModel )
	{
		string fullName =  path +fname+ "." + m_svmNameEXT;
		int result = svm_save_model(fullName.c_str(), m_svmModel);
		if( 0 == result)
		{
			retVal = true; 
		}

	}
	return retVal; 
}
/*******************************************************************/
bool CSVMClassificationModule::SaveClassNameFile( string fname, string path )
{
	bool retVal = false; 
	string fullName =  path + fname + "." + m_classNameEXT;
	ofstream ofs;
	ofs.open(fullName.c_str());
	if( ofs.is_open() )
	{
		int sz=static_cast<int>(m_ClassNameMap.size());
		ofs << sz << endl;
		TClassToNameIter iter; 
		for( iter = m_ClassNameMap.begin(); iter != m_ClassNameMap.end(); ++iter )
		{
			ofs << iter->first << "\t" << iter->second << endl; 
		}
		if( ofs.good() )
		{
			retVal = true;
		}
		ofs.close();
	}
	return retVal; 
}
/*******************************************************************/